
import torch
import torch.nn as nn
import torchvision.models as models


class SPP(nn.Module):
    def __init__(self, M='[1,2,4]'):
        super(SPP, self).__init__()
        self.pooling_4x4 = nn.AdaptiveAvgPool2d((4, 4))
        self.pooling_2x2 = nn.AdaptiveAvgPool2d((2, 2))
        self.pooling_1x1 = nn.AdaptiveAvgPool2d((1, 1))

        self.M = M
        print(self.M)

    def forward(self, x):
        x_4x4 = self.pooling_4x4(x)
        x_2x2 = self.pooling_2x2(x_4x4)
        x_1x1 = self.pooling_1x1(x_4x4)

        x_4x4_flatten = torch.flatten(x_4x4, start_dim=2, end_dim=3)  # B X C X feature_num

        x_2x2_flatten = torch.flatten(x_2x2, start_dim=2, end_dim=3)

        x_1x1_flatten = torch.flatten(x_1x1, start_dim=2, end_dim=3)

        if self.M == '[1,2,4]':
            x_feature = torch.cat((x_1x1_flatten, x_2x2_flatten, x_4x4_flatten), dim=2)
        elif self.M == '[1,2]':
            x_feature = torch.cat((x_1x1_flatten, x_2x2_flatten), dim=2)
        elif self.M=='[1]':
            x_feature = x_1x1_flatten
        else:
            raise NotImplementedError('ERROR M')

        x_strength = x_feature.permute((2, 0, 1))
        x_strength = torch.mean(x_strength, dim=2)


        return x_feature, x_strength

class HookTool:
    def __init__(self):
        self.fea = []

    def hook_fun(self, module, fea_in, fea_out):
        self.fea.append(fea_out)

    def reset(self):
        self.fea = []

class ResNet18_Feature(nn.Module):
    def __init__(self, pretrained=True):
        super(ResNet18_Feature, self).__init__()
        # 加载官方的 ResNet18
        self.resnet18 = models.resnet18(pretrained=pretrained)
        # self.resnet18.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=2, bias=False)
        # self.resnet18.maxpool = nn.Identity()
        # 删除/替换原有的全连接层或保持不变都可以。
        # 这里不需要用到它，所以可以视需要设置为 nn.Identity():
        self.resnet18.fc = nn.Identity()

    def forward(self, x):
        # 与官方 ResNet18 的 forward 对照写法：直接在 forward 里手动调用各层
        x = self.resnet18.conv1(x)
        x = self.resnet18.bn1(x)
        x = self.resnet18.relu(x)
        x = self.resnet18.maxpool(x)

        x = self.resnet18.layer1(x)
        x = self.resnet18.layer2(x)
        x = self.resnet18.layer3(x)
        x = self.resnet18.layer4(x)
        # ------------------------------------------------------------
        # 此时 x 是 (B, 512, H, W)，对于输入 224×224 通常是 (B, 512, 7, 7)
        # ------------------------------------------------------------
        # 我们想要“池化前”特征，所以此时就可以直接返回 x。
        return x


if __name__ == "__main__":
    model = ResNet18_Feature(pretrained=True)
    model.eval()
    spp = SPP()

    cur_hook = HookTool()
    cur_fea = model.resnet18.layer4.register_forward_hook(cur_hook.hook_fun)

    dummy_input = torch.randn(256, 3, 224, 224)  # batch_size=2, 3通道, 224×224
    with torch.no_grad():
        features = model(dummy_input)
        x_feature, x_strength = spp(features)      
        x_spp = x_feature.permute((2, 0, 1))

        m, b, c = x_spp.shape[0], x_spp.shape[1], x_spp.shape[2]
        x_spp = torch.reshape(x_spp, (m * b, c))
    print("x_spp.shape:", x_spp.shape)  # 通常是 (2, 512, 7, 7)
    print("x_feature.shape:", x_feature.shape)  # 通常是 (2, 512, 7, 7)
    print("x_strength.shape:", x_strength.shape)  # 通常是 (2, 512, 7, 7)
    print("features.shape:", features.shape)  # 通常是 (2, 512, 7, 7)



--dataset cifar100 --encoder resnet18 --data_dir /root/cifar100 --seed 5 --split_strategy class --max_epochs 1 --num_tasks 5 --gpus 1 --precision 16 --optimizer sgd --lars --grad_clip_lars --eta_lars 0.02 --exclude_bias_n_norm --scheduler warmup_cosine --lr 1.0 --classifier_lr 0.1 --weight_decay 1e-5 --batch_size 256 --num_workers 8 --brightness 0.4 --contrast 0.4 --saturation 0.2 --hue 0.1 --gaussian_prob 0.0 0.0 --solarization_prob 0.0 0.2 --name byol-sdd-cifar100-5T-cassle --wandb --project UCIL-cifar100 --entity fzl194 --save_checkpoint --method byol_sdd --output_dim 256 --proj_hidden_dim 4096 --pred_hidden_dim 4096 --base_tau_momentum 0.99 --final_tau_momentum 1.0 --momentum_classifier --memory 2000 --checkpoint_dir ./experiments/2025_03_06_01_12_48-byol-sdd-cifar100-5T-cassle --task_idx 1 --offline 